import os
from os.path import join, dirname, abspath
import sys
sys.path.append(join(dirname(abspath(__file__)), ".."))
from methods.utils.utils import calc_world_pts
from matplotlib import animation
from matplotlib.animation import FuncAnimation
import matplotlib.pyplot as plt
import numpy as np
import pickle
import torch
from pathlib import Path
import matplotlib.image as mpimg
from matplotlib.colors import to_rgba
sys.path.append(join(dirname(abspath(__file__)), "../../.."))
from utils.parse_object import load_obj_mesh


connections = [(0, 1), (1, 2), (2, 3), (3, 4), (0, 5), (5, 6), (6, 7), (7, 8), (0, 9), (9, 10), (10, 11), (11, 12), (0, 13), (13, 14), (14, 15), (15, 16), (0, 17), (17, 18), (18, 19), (19, 20)]  # MANO hand

obj_connects = {}

def visualize_skeleton(skeledonData, objPoint, save_dir = './test.gif'):
    fig = plt.figure()
    ax = fig.add_subplot(111 , projection="3d")
    def updateHumanObj(frame, *fargs):

        ax.clear()
        bodyData, objPoint, scat = fargs
        z_points = bodyData[frame, :, 2] #* -1.0
        x_points = bodyData[frame, :, 0]
        y_points = bodyData[frame, :, 1]
        for connect in connections:
            a,b = connect
            ax.plot([x_points[a], x_points[b]],[y_points[a],y_points[b]],[z_points[a],z_points[b]], color="b")
        ax.scatter3D(x_points, y_points, z_points, color="r")

        thisObjPoint = objPoint[frame].reshape((-1,3))
        z_points = thisObjPoint[ :, 2] #* -1.0
        x_points = thisObjPoint[ :, 0] 
        y_points = thisObjPoint[ :, 1]
        ax.scatter3D(x_points, y_points, z_points, color="g")
        #'''
        ax.plot([0.0, 0.0],[-0.7,1.2,],[0.0,0.0], color="w")#b
        ax.plot([-1.2,1.2],[0.0,0.0,],[0.0,0.0], color="w")#r
        ax.plot([0.0, 0.0],[0.0,0.0,],[-1.2,1.2], color="w")#g
        # plt.axis('off')
        ax.set_axis_off()
        return ax


    def visualize(skeledonData, objPoint, startFrame = 0):
        bodyData = skeledonData
        lenFrame = skeledonData.shape[0]
        
        bodyData = bodyData.reshape((lenFrame, 21, 3))
        # ax.yaxis.set_label_position("top")
        ax.view_init(elev=117., azim=-88.)
        scat = ax.scatter(bodyData[0,:,0], bodyData[0,:,1], bodyData[0,:,2], c='r', marker = 'o',alpha=0.5, s=100)
        
        #time.sleep(.01)
        ani = animation.FuncAnimation(fig, updateHumanObj, frames= range(lenFrame), interval = 50, repeat_delay=100,
                                    fargs=(bodyData, objPoint, scat))
        ani.save(save_dir)

    visualize(skeledonData, objPoint)
    plt.clf()
    plt.close()


def visualize_skeleton_pred_gt(skeledonData, objPoint, skeledonData_gt, objPoint_gt, save_dir = './test.gif', elev=117., azim=-88., roll=0.0, obj_name='chair4', tmp_dir='./tmp'):
    fig = plt.figure()
    ax = fig.add_subplot(111 , projection="3d")
    
    Path(tmp_dir).mkdir(exist_ok=True, parents=True)
    ax.view_init(elev=elev, azim=azim, roll=roll,vertical_axis='y')
    def updateHumanObj(frame, *fargs):

        ax.clear()
        ax.plot([0.0, 0.0],[-0.7,1.2,],[0.0,0.0], color="w")#b
        ax.plot([-1.2,1.2],[0.0,0.0,],[0.0,0.0], color="w")#r
        ax.plot([0.0, 0.0],[0.0,0.0,],[-1.2,1.2], color="w")#g
        # ax.plot([0.0, 0.0],[-0.49,0.7,],[0.0,0.0], color="w")#b
        # ax.plot([-0.7,0.7],[0.0,0.0,],[0.0,0.0], color="w")#r
        # ax.plot([0.0, 0.0],[0.0,0.0,],[-0.7,0.7], color="w")#g
        
        bodyData, objPoint, bodyData_gt, objPoint_gt = fargs
        obj_connections = obj_connects.get(obj_name,[])
        if frame>10:
            # visualize pred
            z_points = bodyData[frame, :, 2] #* -1.0
            x_points = bodyData[frame, :, 0]
            y_points = bodyData[frame, :, 1]
            p_human_color = to_rgba("seagreen")
            p_human_color = (*p_human_color[:3],0.8)
            for connect in connections:
                a,b = connect
                ax.plot([x_points[a], x_points[b]],[y_points[a],y_points[b]],[z_points[a],z_points[b]], color=p_human_color)
            # ax.scatter3D(x_points, y_points, z_points, color="purple")

            thisObjPoint = objPoint[frame].reshape((12,3))
            z_points = thisObjPoint[ :, 2] #* -1.0
            x_points = thisObjPoint[ :, 0] 
            y_points = thisObjPoint[ :, 1]
            
            p_obj_color = to_rgba("salmon")
            p_obj_color = (*p_obj_color[:3],0.8)
            for connect in obj_connections:
                a,b = connect
                ax.plot([x_points[a], x_points[b]],[y_points[a],y_points[b]],[z_points[a],z_points[b]], color=p_obj_color)
            # ax.scatter3D(x_points, y_points, z_points, color="r")

        # NOTE: plot gt
        if frame<objPoint_gt.shape[0]:
            gt_human_color = to_rgba("lightgrey")
            gt_human_color = (*gt_human_color[:3],0.8)
            z_points_gt = bodyData_gt[frame, :, 2] #* -1.0
            x_points_gt = bodyData_gt[frame, :, 0]
            y_points_gt = bodyData_gt[frame, :, 1]
            for connect in connections:
                a,b = connect
                ax.plot([x_points_gt[a], x_points_gt[b]],[y_points_gt[a],y_points_gt[b]],[z_points_gt[a],z_points_gt[b]], color=gt_human_color)#royalblue 
            # ax.scatter3D(x_points_gt, y_points_gt, z_points_gt, color="b")

            thisObjPoint = objPoint_gt[frame].reshape((12,3))
            
            z_points_gt = thisObjPoint[ :, 2] #* -1.0
            x_points_gt = thisObjPoint[ :, 0] 
            y_points_gt = thisObjPoint[ :, 1]
            gt_obj_color = to_rgba("dimgrey")
            gt_obj_color = (*gt_obj_color[:3],0.5)
            for connect in obj_connections:
                a,b = connect
                ax.plot([x_points_gt[a], x_points_gt[b]],[y_points_gt[a],y_points_gt[b]],[z_points_gt[a],z_points_gt[b]], color=gt_obj_color)#teal
        # ax.scatter3D(x_points, y_points, z_points, color="g")

        #'''

        ax.set_axis_off()
        return ax


    def visualize(skeledonData, objPoint, skeledonData_gt, objPoint_gt, startFrame = 0):
        real_save_dir = save_dir[:-4]+'.gif'
        bodyData = skeledonData
        lenFrame = skeledonData.shape[0]
        ax.view_init(elev=elev, azim=azim, roll=roll,vertical_axis='y')
        # gt: blue and green
        # pred: purple and red
        
        bodyData = bodyData.reshape((lenFrame, 21, 3))
        scat = ax.scatter(bodyData[0,:,0], bodyData[0,:,1], bodyData[0,:,2], c='b', marker = 'o',alpha=0.5, s=100)
        
        #time.sleep(.01)
        ani = animation.FuncAnimation(fig, updateHumanObj, frames= range(lenFrame), interval = 50, repeat_delay=100,
                                    fargs=(bodyData, objPoint, skeledonData_gt, objPoint_gt))
        ani.save(real_save_dir)
        
    def visualize_imgs(skeledonData, objPoint, skeledonData_gt, objPoint_gt):
        real_save_dir = save_dir[:-4]+'.png'
        total = 20
        steps = [3,7,11,15,19]
        images = []
        x_scale = 779
        y_scale = 779#dpi=200
        x_tight_amnt_left = 60*2
        x_tight_amnt_right = 60*2
        x_rescaled = x_scale - x_tight_amnt_left - x_tight_amnt_right
        y_tight_amnt_upper = 60*2
        y_tight_amnt_lower = 100*2
        y_rescaled = y_scale - y_tight_amnt_upper  - y_tight_amnt_lower
        
        
        for frame in steps:
            fargs=(skeledonData, objPoint,  skeledonData_gt, objPoint_gt)
            ax = updateHumanObj(frame, *fargs)
            tmppath = str(Path(tmp_dir)/f"{frame}.png")
            plt.savefig(tmppath, bbox_inches='tight',dpi=200)
            # print(mpimg.imread(tmppath).shape)
            images.append(mpimg.imread(tmppath)[y_tight_amnt_upper:y_tight_amnt_upper+y_rescaled,x_tight_amnt_left:x_tight_amnt_left+x_rescaled])
            # 480,640,3
            # 0-255

        # concat imgs
        big_img = np.zeros((y_rescaled, (x_rescaled*5),4))
        for i in range(len(steps)):
            big_img[:,x_rescaled*i:x_rescaled*(i+1)]=images[i]
        mpimg.imsave(real_save_dir, big_img)
        
    
        # output, an array of images at different timesteps
    # visualize_imgs(skeledonData, objPoint, skeledonData_gt, objPoint_gt)
    visualize(skeledonData, objPoint, skeledonData_gt, objPoint_gt)
    plt.clf()
    plt.close()
    

# liuyun
def visualize_4instances_old(body_pred, body_gt, obj_pred, obj_gt, batch_idx=0, save_dir=None, elev=90.0, azim=180.0, roll=0.0, vertical_axis='z'):
    # 放大的倍数
    alpha = 10.0
    # visualize the first clip in the batch
    T = body_pred.shape[0]
    
    # skeledonData = body_pred[:, 0].cpu().numpy().reshape(T, 2, 21, 3) * 10.0
    # objData = obj_pred[:, 0].cpu().numpy().reshape(T, 2, -1, 3) * 10.0
    # skeledonData_gt = body_gt[:, 0].cpu().numpy().reshape(T, 2, 21, 3) * 10.0
    # objData_gt = obj_gt[:, 0].cpu().numpy().reshape(T, 2, -1, 3) * 10.0

    if save_dir is None:
        save_dir = "./visualization/chmp_train"
    os.makedirs(save_dir, exist_ok=True)
    rend_video_path = join(save_dir, '{}_{}_{}_all.gif'.format("chmp", batch_idx, 0))

    fig = plt.figure()
    ax = fig.add_subplot(111 , projection="3d")
    ax.view_init(elev=elev, azim=azim, roll=roll, vertical_axis=vertical_axis)
    
    def updateHumanObj(frame, *fargs):
        ax.clear()
        ax.plot([0.0, 0.0],[-0.7, 1.2],[0.0,0.0], color="w")#b
        ax.plot([-1.2, 1.2],[0.0, 0.0],[0.0,0.0], color="w")#r
        ax.plot([0.0, 0.0],[0.0, 0.0],[-1.2,1.2], color="w")#g
        
        bodyData, objPoint, bodyData_gt, objPoint_gt = fargs
        
        if frame > 10:
            # visualize pred
            z_points = bodyData[frame, 0, :, 2] * alpha
            x_points = bodyData[frame, 0, :, 0] * alpha
            y_points = bodyData[frame, 0, :, 1] * alpha
            p_human_color = to_rgba("seagreen")
            p_human_color = (*p_human_color[:3], 0.8)
            for connect in connections:
                a, b = connect
                ax.plot([x_points[a], x_points[b]], [y_points[a], y_points[b]], [z_points[a], z_points[b]], color=p_human_color)
            ax.scatter3D(x_points, y_points, z_points, color="purple")
            z_points = bodyData[frame, 1, :, 2] * alpha
            x_points = bodyData[frame, 1, :, 0] * alpha
            y_points = bodyData[frame, 1, :, 1] * alpha
            p_human_color = to_rgba("seagreen")
            p_human_color = (*p_human_color[:3], 0.8)
            for connect in connections:
                a, b = connect
                ax.plot([x_points[a], x_points[b]], [y_points[a], y_points[b]], [z_points[a], z_points[b]], color=p_human_color)
            ax.scatter3D(x_points, y_points, z_points, color="purple")

            thisObjPoint = objPoint[frame][0].reshape((-1, 3)) * alpha
            z_points = thisObjPoint[:, 2]
            x_points = thisObjPoint[:, 0]
            y_points = thisObjPoint[:, 1]
            p_obj_color = to_rgba("salmon")
            p_obj_color = (*p_obj_color[:3], 0.8)
            ax.scatter3D(x_points, y_points, z_points, color="r")
            thisObjPoint = objPoint[frame][1].reshape((-1, 3)) * alpha
            z_points = thisObjPoint[:, 2]
            x_points = thisObjPoint[:, 0]
            y_points = thisObjPoint[:, 1]
            p_obj_color = to_rgba("salmon")
            p_obj_color = (*p_obj_color[:3], 0.8)
            ax.scatter3D(x_points, y_points, z_points, color="r")

        # NOTE: plot gt
        if frame < T:
            z_points_gt = bodyData_gt[frame, 0, :, 2] * alpha
            x_points_gt = bodyData_gt[frame, 0, :, 0] * alpha
            y_points_gt = bodyData_gt[frame, 0, :, 1] * alpha
            gt_human_color = to_rgba("lightgrey")
            gt_human_color = (*gt_human_color[:3], 0.8)
            for connect in connections:
                a,b = connect
                ax.plot([x_points_gt[a], x_points_gt[b]], [y_points_gt[a], y_points_gt[b]], [z_points_gt[a], z_points_gt[b]], color=gt_human_color) #royalblue 
            ax.scatter3D(x_points_gt, y_points_gt, z_points_gt, color="b")
            z_points_gt = bodyData_gt[frame, 1, :, 2] * alpha
            x_points_gt = bodyData_gt[frame, 1, :, 0] * alpha
            y_points_gt = bodyData_gt[frame, 1, :, 1] * alpha
            gt_human_color = to_rgba("lightgrey")
            gt_human_color = (*gt_human_color[:3], 0.8)
            for connect in connections:
                a,b = connect
                ax.plot([x_points_gt[a], x_points_gt[b]], [y_points_gt[a], y_points_gt[b]], [z_points_gt[a], z_points_gt[b]], color=gt_human_color) #royalblue 
            ax.scatter3D(x_points_gt, y_points_gt, z_points_gt, color="b")

            thisObjPoint = objPoint_gt[frame][0].reshape((-1, 3)) * alpha
            z_points_gt = thisObjPoint[:, 2]
            x_points_gt = thisObjPoint[:, 0]
            y_points_gt = thisObjPoint[:, 1]
            gt_obj_color = to_rgba("dimgrey")
            gt_obj_color = (*gt_obj_color[:3], 0.5)
            ax.scatter3D(x_points_gt, y_points_gt, z_points_gt, color="g")
            thisObjPoint = objPoint_gt[frame][1].reshape((-1, 3)) * alpha
            z_points_gt = thisObjPoint[:, 2]
            x_points_gt = thisObjPoint[:, 0]
            y_points_gt = thisObjPoint[:, 1]
            gt_obj_color = to_rgba("dimgrey")
            gt_obj_color = (*gt_obj_color[:3], 0.5)
            ax.scatter3D(x_points_gt, y_points_gt, z_points_gt, color="g")

        ax.set_axis_off()
        return ax

    def visualize(skeledonData, objPoint, skeledonData_gt, objPoint_gt, startFrame = 0):
        # gt: blue and green
        # pred: purple and red
        
        # time.sleep(0.01)
        ani = animation.FuncAnimation(fig, updateHumanObj, frames=range(T), interval=30, repeat_delay=100, fargs=(skeledonData, objPoint, skeledonData_gt, objPoint_gt))
        ani.save(rend_video_path)
    
    visualize(body_pred, obj_pred, body_gt, obj_gt)
    plt.clf()
    plt.close()


def visualize_4instances(body_pred, body_gt, obj_pred, obj_gt, pose_pred, pose_gt, sequence_name, obj_names, batch_idx=0, save_dir=None, elev=90.0, azim=180.0, roll=0.0, vertical_axis='z'):
    """
    pose: [px, py, pz, x, y, z, w]
    """
    
    # 只可视化第0个clip
    T = body_pred.shape[0]
    body_pred = body_pred[:, 0]  # (T, 2, 21, 3)
    body_gt = body_gt[:, 0]  # (T, 2, 21, 3)
    pose_pred = pose_pred[:, 0]  # (T, 2, 7)
    pose_gt = pose_gt[:, 0]  # (T, 2, 7)
    
    tool_name, target_name = obj_names[0][0], obj_names[1][0]
    obj_model_dir = "/data2/HOI-mocap/object_models_final_simplied"
    device = body_pred.device
    tool_model_pts = torch.from_numpy(load_obj_mesh(join(obj_model_dir, tool_name + "_cm.obj"), unit=0.01).vertices).float().to(device)  # (N1, 3)
    target_model_pts = torch.from_numpy(load_obj_mesh(join(obj_model_dir, target_name + "_cm.obj"), unit=0.01).vertices).float().to(device)  # (N2, 3)
    tool_pred_pts = calc_world_pts(pose_pred[:, 0], tool_model_pts)  # (T, N1, 3)
    tool_gt_pts = calc_world_pts(pose_gt[:, 0], tool_model_pts)  # (T, N1, 3)
    target_pred_pts = calc_world_pts(pose_pred[:, 1], target_model_pts)  # (T, N2, 3)
    target_gt_pts = calc_world_pts(pose_gt[:, 1], target_model_pts)  # (T, N2, 3)
    
    obj_pred = []
    for i in range(T):
        obj_pred.append([tool_pred_pts[i].detach().cpu().numpy(), target_pred_pts[i].detach().cpu().numpy()])
    obj_gt = []
    for i in range(T):
        obj_gt.append([tool_gt_pts[i].detach().cpu().numpy(), target_gt_pts[i].detach().cpu().numpy()])
    
    visualize_4instances_old(body_pred.detach().cpu().numpy(), body_gt.detach().cpu().numpy(), obj_pred, obj_gt, batch_idx=batch_idx, save_dir=save_dir, elev=elev, azim=azim, roll=roll, vertical_axis=vertical_axis)
